"""
分割与重建的第四步，将vtp表面模型转换为glb文件
这里需要优化，后端只做统一的建模，不负责颜色设置，而是交给前端来设置
"""
import os
import vtk

def get_organ_properties(organ_name: str) -> dict:
    properties = {
        'aorta': {'color': (1.0, 0.0, 0.0), 'opacity': 0.9},
        'liver': {'color': (0.7, 0.2, 0.2), 'opacity': 0.7},
        'spleen': {'color': (0.6, 0.2, 0.8), 'opacity': 0.7},
        'stomach': {'color': (1.0, 0.5, 0.8), 'opacity': 0.6},
        'kidney_right': {'color': (0.9, 0.7, 0.1), 'opacity': 0.8},
        'kidney_left': {'color': (0.9, 0.7, 0.1), 'opacity': 0.8},
        'portal_vein_and_splenic_vein': {'color': (0.0, 0.0, 1.0), 'opacity': 0.9},
        'inferior_vena_cava': {'color': (0.0, 0.0, 0.7), 'opacity': 0.9},
        'pancreas': {'color': (1.0, 1.0, 0.6), 'opacity': 0.8},
        'gallbladder': {'color': (0.0, 0.7, 0.0), 'opacity': 0.8},
        'default_bone': {'color': (0.94, 0.92, 0.84), 'opacity': 1.0},
        'default_color': {'color': (0.8, 0.8, 0.8), 'opacity': 0.5}
    }
    if organ_name in properties: return properties[organ_name]
    bone_keywords = ['rib', 'vertebrae', 'femur', 'hip', 'humerus', 'scapula', 'clavicula', 'sternum', 'skull', 'sacrum']
    if any(keyword in organ_name for keyword in bone_keywords): return properties['default_bone']
    return properties['default_color']


def convert_all_vtps_to_glbs(vtp_input_dir: str, glb_output_dir: str):
    """
    遍历VTP目录，为每个VTP文件生成一个独立的、着色后的GLB文件。
    """
    print(f"--- 开始批量转换 VTP -> GLB ---")
    print(f"VTP 输入源: {vtp_input_dir}")
    print(f"GLB 输出目标: {glb_output_dir}")

    os.makedirs(glb_output_dir, exist_ok=True)
    vtp_files = [f for f in os.listdir(vtp_input_dir) if f.endswith('.vtp')]

    if not vtp_files:
        print(f"错误: 目录 {vtp_input_dir} 中未找到任何.vtp文件。")
        return

    for vtp_filename in vtp_files:
        vtp_path = os.path.join(vtp_input_dir, vtp_filename)
        organ_name = vtp_filename.split('.')[0]
        glb_path = os.path.join(glb_output_dir, f"{organ_name}.glb")
        
        print(f"  - 正在转换: {vtp_filename} -> {os.path.basename(glb_path)}")
        
        # 读取VTP
        reader = vtk.vtkXMLPolyDataReader()
        reader.SetFileName(vtp_path)
        reader.Update()

        # 创建Actor并设置属性
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputConnection(reader.GetOutputPort())
        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        props = get_organ_properties(organ_name)
        actor.GetProperty().SetColor(props['color'])
        actor.GetProperty().SetOpacity(props['opacity'])
        
        # 创建独立的渲染场景用于导出
        renderer = vtk.vtkRenderer()
        render_window = vtk.vtkRenderWindow()
        render_window.AddRenderer(renderer)
        renderer.AddActor(actor)
        
        # 导出为GLB
        exporter = vtk.vtkGLTFExporter()
        exporter.SetRenderWindow(render_window)
        exporter.SetFileName(glb_path)
        exporter.InlineDataOn()
        exporter.Write()

    print(f"\n🎉 成功！已在输出目录中生成 {len(vtp_files)} 个GLB模型文件。")


if __name__ == '__main__':
    BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    VTP_MODELS_DIR = os.path.join(BASE_DIR, r"data\\vtp_models\\patient001\Abdomen-V  2.0  B30f\\all_organs")
    GLB_OUTPUT_DIR = os.path.join(BASE_DIR, r"static\\glb\\patient001\Abdomen-V  2.0  B30f\\all_organs_glbs")

    convert_all_vtps_to_glbs(VTP_MODELS_DIR, GLB_OUTPUT_DIR)